#!/usr/bin/env python3
"""
高级训练额外参数节点及参数合并节点
采样参数、正则化参数分别独立成节点
"""

from .lora_trainer_utils.constants import CHARACTER_LORA_PARAMS_CATEGORY
import os

class AdvancedSamplingParamsNode:
    """高级采样参数节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "sample_every_n_steps": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 10000,
                    "tooltip": "每多少步采样输出一次样本图像，0为不采样"
                }),
                "sample_every_n_epochs": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100,
                    "tooltip": "每多少epoch采样输出一次样本图像，0为不采样"
                }),
                "sample_prompts": ("STRING", {
                    "default": "",
                    "tooltip": "采样用prompt文件路径（.txt/.json/.toml格式）",
                    "multiline": True
                }),
                "sample_sampler": (["ddim", "pndm", "heun", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a"], {
                    "default": "k_euler_a",
                    "tooltip": "采样器算法"
                }),
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数 - 只验证路径是否存在"""
        # 检查采样提示词文件（如果提供）
        sample_prompts = kwargs.get("sample_prompts")
        if sample_prompts and isinstance(sample_prompts, str) and sample_prompts.strip() and not os.path.exists(sample_prompts):
            return f"采样提示词文件不存在: {sample_prompts}"
        
        return True
    
    RETURN_TYPES = ("TRAINING_EXTRA_PARAMS",)
    RETURN_NAMES = ("采样参数",)
    FUNCTION = "generate_sampling_params"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY
    def generate_sampling_params(self, **kwargs):
       
        return (kwargs,)

class AdvancedRegularizationParamsNode:
    """高级正则化参数节点"""
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "reg_data_dir": ("STRING", {
                    "default": "",
                    "tooltip": "正则化图像文件夹（DreamBooth专用）"
                }),
                "prior_loss_weight": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 10.0,
                    "tooltip": "正则化损失权重"
                }),
            }
        }
    RETURN_TYPES = ("TRAINING_EXTRA_PARAMS",)
    RETURN_NAMES = ("正则化参数",)
    FUNCTION = "generate_regularization_params"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY
    def generate_regularization_params(self, **kwargs):
        return (kwargs,)

class AdvancedTrainingExtraParamsNode:
    """高级训练额外参数节点（保留：分桶、caption扩展名、文本编码器训练控制、分词器填充等）"""
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # caption扩展名
                "caption_extension": ("STRING", {
                    "default": ".txt",
                    "tooltip": "txt文件扩展名（如需用.caption可修改）"
                }),
                # 分桶相关
                "enable_bucket": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "是否启用分桶（Aspect Ratio Bucketing）"
                }),
                "min_bucket_reso": ("INT", {
                    "default": 256,
                    "min": 64,
                    "max": 2048,
                    "tooltip": "分桶最小分辨率"
                }),
                "max_bucket_reso": ("INT", {
                    "default": 1024,
                    "min": 64,
                    "max": 4096,
                    "tooltip": "分桶最大分辨率"
                }),
                "bucket_reso_steps": ("INT", {
                    "default": 64,
                    "min": 1,
                    "max": 512,
                    "tooltip": "分桶分辨率步长"
                }),
                "bucket_no_upscale": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "分桶时不放大图片"
                }),
                # 文本编码器训练控制
                "stop_text_encoder_training": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100000,
                    "tooltip": "在指定步数后停止训练文本编码器，0为不停止"
                }),
                # 分词器填充
                "no_token_padding": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "不进行分词器填充"
                }),
            }
        }
    RETURN_TYPES = ("TRAINING_EXTRA_PARAMS",)
    RETURN_NAMES = ("高级训练额外参数",)
    FUNCTION = "generate_extra_params"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY
    def generate_extra_params(self, **kwargs):
        return (kwargs,)

class AdvancedExtraParamsMergerNode:
    """高级训练参数合并节点：合并主训练参数和多个额外参数（支持多个独立输入）"""
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "main_params": ("TRAINING_PARAMS", {"tooltip": "主训练参数（如角色参数节点输出）"}),
            },
            "optional": {
                "extra_params1": ("TRAINING_EXTRA_PARAMS", {"tooltip": "额外参数1"}),
                "extra_params2": ("TRAINING_EXTRA_PARAMS", {"tooltip": "额外参数2"}),
                "extra_params3": ("TRAINING_EXTRA_PARAMS", {"tooltip": "额外参数3"}),
                "extra_params4": ("TRAINING_EXTRA_PARAMS", {"tooltip": "额外参数4"}),
            }
        }
    RETURN_TYPES = ("TRAINING_PARAMS",)
    RETURN_NAMES = ("合并后训练参数",)
    FUNCTION = "merge_params"
    CATEGORY = CHARACTER_LORA_PARAMS_CATEGORY

    def merge_params(self, main_params, extra_params1=None, extra_params2=None, extra_params3=None, extra_params4=None):
        merged = dict(main_params)
        for extra in [extra_params1, extra_params2, extra_params3, extra_params4]:
            if extra:
                merged.update(extra)
        return (merged,)

NODE_CLASS_MAPPINGS = {
    "AdvancedSamplingParamsNode": AdvancedSamplingParamsNode,
    "AdvancedRegularizationParamsNode": AdvancedRegularizationParamsNode,
    "AdvancedTrainingExtraParamsNode": AdvancedTrainingExtraParamsNode,
    "AdvancedExtraParamsMergerNode": AdvancedExtraParamsMergerNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "AdvancedSamplingParamsNode": "高级采样参数",
    "AdvancedRegularizationParamsNode": "高级正则化参数",
    "AdvancedTrainingExtraParamsNode": "高级训练额外参数",
    "AdvancedExtraParamsMergerNode": "高级参数合并",
}
